//go:build windows
// +build windows

package net

import (
	"fmt"
	"go-cmd/internal/common"
	"go-cmd/internal/utils"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

const anySize = 1

// the status of tcp conntion
type TCP_STATE uint32

//go:generate stringer -type=TCP_STATE
const (
	MIB_TCP_STATE_CLOSED TCP_STATE = iota + 1
	MIB_TCP_STATE_LISTEN
	MIB_TCP_STATE_SYN_SENT
	MIB_TCP_STATE_SYN_RCVD
	MIB_TCP_STATE_ESTAB
	MIB_TCP_STATE_FIN_WAIT1
	MIB_TCP_STATE_FIN_WAIT2
	MIB_TCP_STATE_CLOSE_WAIT
	MIB_TCP_STATE_CLOSING
	MIB_TCP_STATE_LAST_ACK
	MIB_TCP_STATE_TIME_WAIT
	MIB_TCP_STATE_DELETE_TCB
)

type TCP_TABLE_CLASS int32

const (
	TCP_TABLE_BASIC_LISTENER = iota
	TCP_TABLE_BASIC_CONNECTIONS
	TCP_TABLE_BASIC_ALL
	TCP_TABLE_OWNER_PID_LISTENER
	TCP_TABLE_OWNER_PID_CONNECTIONS
	TCP_TABLE_OWNER_PID_ALL
	TCP_TABLE_OWNER_MODULE_LISTENER
	TCP_TABLE_OWNER_MODULE_CONNECTIONS
	TCP_TABLE_OWNER_MODULE_ALL
)

var (
	modIpHelpApi = common.ModIphlpApi
	// See more info at https://learn.microsoft.com/zh-cn/windows/win32/api/iphlpapi/nf-iphlpapi-getextendedtcptable
	procGetTCPTable = modIpHelpApi.NewProc("GetExtendedTcpTable")
	procGetUDPTable = modIpHelpApi.NewProc("GetExtendedUdpTable")
	procGetIfEntry  = modIpHelpApi.NewProc("GetIfEntry")
)

// See more info at https://learn.microsoft.com/zh-cn/windows/win32/api/tcpmib/ns-tcpmib-mib_tcprow_lh
// more info at https://learn.microsoft.com/zh-cn/windows/win32/api/tcpmib/ns-tcpmib-mib_tcptable
type MIB_TCPROW struct {
	DwState      uint32
	State        uint32
	DwLocalAddr  uint32
	DwLocalPort  uint32
	DwRemoteAddr uint32
	DwRemotePort uint32
}
type MIB_TCPTABLE struct {
	DwNumEntries uint32
	Table        [anySize]MIB_TCPROW
}
type MIB_TCP6ROW struct {
	State           uint32
	LocalAddr       uint64
	DwLocalScopeId  uint32
	DwLocalPort     uint32
	RemoteAddr      uint64
	DwRemoteScopeId uint32
	DwRemotePort    uint32
}
type MIB_TCP6TABLE struct {
	DwNumEntries uint32
	Table        [anySize]MIB_TCP6ROW
}

type MIB_TCPROW_OWNER_PID struct {
	DwState      TCP_STATE
	DwLocalAddr  uint32
	DwLocalPort  uint32
	DwRemoteAddr uint32
	DwRemotePort uint32
	DwOwningPid  uint32
}
type MIB_TCPTABLE_OWNER_PID struct {
	DwNumEntries uint32
	Table        [anySize]MIB_TCPROW_OWNER_PID
}

func (m *MIB_TCPROW_OWNER_PID) convertToConnection() ConnectionStat {
	return ConnectionStat{
		Family: kindTCP4.family,
		Type:   kindTCP4.sockType,
		Laddr: Addr{
			IP:   utils.ParseIPv4HexString(m.DwLocalAddr),
			Port: uint32(windows.Ntohs(uint16(m.DwLocalPort))),
		},
		Raddr: Addr{
			IP:   utils.ParseIPv4HexString(m.DwRemoteAddr),
			Port: uint32(windows.Ntohs(uint16(m.DwRemotePort))),
		},
		Pid:    int(m.DwOwningPid),
		Status: m.DwState.String(),
	}
}

type MIB_TCP6ROW_OWNER_PID struct {
	UcLocalAddr     [16]byte
	DwLocalScopeId  uint32
	DwLocalPort     uint32
	UcRemoteAddr    [16]byte
	DwRemoteScopeId uint32
	DwRemotePort    uint32
	DwState         TCP_STATE
	DwOwningPid     uint32
}
type MIB_TCP6TABLE_OWNER_PID struct {
	DwNumEntries uint32
	Table        [anySize]MIB_TCP6ROW_OWNER_PID
}

func (m *MIB_TCP6ROW_OWNER_PID) convertToConnection() ConnectionStat {
	return ConnectionStat{
		Family: kindTCP4.family,
		Type:   kindTCP4.sockType,
		Laddr: Addr{
			IP:   utils.ParseIPv6HexString(m.UcLocalAddr),
			Port: uint32(windows.Ntohs(uint16(m.DwLocalPort))),
		},
		Raddr: Addr{
			IP:   utils.ParseIPv6HexString(m.UcRemoteAddr),
			Port: uint32(windows.Ntohs(uint16(m.DwRemotePort))),
		},
		Pid:    int(m.DwOwningPid),
		Status: m.DwState.String(),
	}
}

type NetConnectionKindType struct {
	family   uint32
	sockType uint32
	fileName string
}

var (
	kindTCP4 = NetConnectionKindType{
		family:   windows.AF_INET,     // ipv4
		sockType: windows.SOCK_STREAM, // byte flow
		fileName: "tcp",
	}
	kindTCP6 = NetConnectionKindType{
		family:   windows.AF_INET6,    // ipv6
		sockType: windows.SOCK_STREAM, // byte flow
		fileName: "tcp6",
	}
	kindUDP4 = NetConnectionKindType{
		family:   windows.AF_INET,    // ipv4
		sockType: windows.SOCK_DGRAM, // packet
		fileName: "udp",
	}
	kindUDP6 = NetConnectionKindType{
		family:   windows.AF_INET6,   // ipv6
		sockType: windows.SOCK_DGRAM, // packet
		fileName: "udp6",
	}
	netConnectionTypeMap = map[string][]NetConnectionKindType{
		"all":   {kindTCP4, kindTCP6, kindUDP4, kindUDP6},
		"tcp":   {kindTCP4, kindTCP6},
		"tcp4":  {kindTCP4},
		"tcp6":  {kindTCP6},
		"udp":   {kindUDP4, kindUDP6},
		"udp4":  {kindUDP4},
		"udp6":  {kindUDP6},
		"inet":  {kindTCP4, kindTCP6, kindUDP4, kindUDP6},
		"inet4": {kindTCP4, kindUDP4},
		"inet6": {kindTCP6, kindUDP6},
	}
)

func GetTCPConnections(family uint32) ([]ConnectionStat, error) {
	var (
		p    uintptr
		buf  []byte
		size uint32

		pmibTCPTable  *MIB_TCPTABLE_OWNER_PID
		pmibTCP6Table *MIB_TCP6TABLE_OWNER_PID
	)
	if family == 0 {
		return nil, fmt.Errorf("family must be required")
	}
	for {
		switch family {
		case kindTCP4.family:
			if len(buf) > 0 {
				pmibTCPTable = (*MIB_TCPTABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
				p = uintptr(unsafe.Pointer(pmibTCPTable))
			} else {
				p = uintptr(unsafe.Pointer(pmibTCPTable))
			}
		case kindTCP6.family:
			if len(buf) > 0 {
				pmibTCP6Table = (*MIB_TCP6TABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
				p = uintptr(unsafe.Pointer(pmibTCP6Table))
			} else {
				p = uintptr(unsafe.Pointer(pmibTCP6Table))
			}
		}
		r, _, _ := procGetTCPTable.Call(p, uintptr(unsafe.Pointer(&size)), 1, uintptr(family), uintptr(TCP_TABLE_OWNER_PID_ALL), uintptr(0))
		var err error
		if r != 0 {
			err = syscall.Errno(r)
		}
		if err == nil {
			break
		}
		if err != windows.ERROR_INSUFFICIENT_BUFFER {
			return nil, err
		}
		buf = make([]byte, size)
	}

	var (
		index, step int
		length      int
	)

	stats := make([]ConnectionStat, 0)
	switch family {
	case kindTCP4.family:
		index, step, length = getTableInfo(kindTCP4.fileName, pmibTCPTable)
	case kindTCP6.family:
		index, step, length = getTableInfo(kindTCP6.fileName, pmibTCP6Table)
	}

	if length == 0 {
		return nil, nil
	}
	for i := 0; i < length; i++ {
		switch family {
		case kindTCP4.family:
			mibs := (*MIB_TCPROW_OWNER_PID)(unsafe.Pointer(&buf[index]))
			ns := mibs.convertToConnection()
			stats = append(stats, ns)
		case kindTCP6.family:
			mibs := (*MIB_TCP6ROW_OWNER_PID)(unsafe.Pointer(&buf[index]))
			ns := mibs.convertToConnection()
			stats = append(stats, ns)
		}
		index += step
	}
	return stats, nil
}

func getTableInfo(filename string, table any) (index, step, length int) {
	switch filename {
	case kindTCP4.fileName:
		index = int(unsafe.Sizeof(table.(*MIB_TCPTABLE_OWNER_PID).DwNumEntries))
		step = int(unsafe.Sizeof(table.(*MIB_TCPTABLE_OWNER_PID).Table))
		length = int(table.(*MIB_TCPTABLE_OWNER_PID).DwNumEntries)
	case kindTCP6.fileName:
		index = int(unsafe.Sizeof(table.(*MIB_TCP6TABLE_OWNER_PID).DwNumEntries))
		step = int(unsafe.Sizeof(table.(*MIB_TCP6TABLE_OWNER_PID).Table))
		length = int(table.(*MIB_TCP6TABLE_OWNER_PID).DwNumEntries)
		// todo: waiting to implement udp and so on
	}
	return
}
